// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2012 Désiré Nuentsa-Wakam <desire.nuentsa_wakam@inria.fr>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

/*

 * NOTE: This file is the modified version of sp_coletree.c file in SuperLU

 * -- SuperLU routine (version 3.1) --
 * Univ. of California Berkeley, Xerox Palo Alto Research Center,
 * and Lawrence Berkeley National Lab.
 * August 1, 2008
 *
 * Copyright (c) 1994 by Xerox Corporation.  All rights reserved.
 *
 * THIS MATERIAL IS PROVIDED AS IS, WITH ABSOLUTELY NO WARRANTY
 * EXPRESSED OR IMPLIED.  ANY USE IS AT YOUR OWN RISK.
 *
 * Permission is hereby granted to use or copy this program for any
 * purpose, provided the above notices are retained on all copies.
 * Permission to modify the code and to distribute modified code is
 * granted, provided the above notices are retained, and a notice that
 * the code was modified is included with the above copyright notice.
 */
#ifndef SPARSE_COLETREE_H
#define SPARSE_COLETREE_H

namespace Eigen {

namespace internal {

/** Find the root of the tree/set containing the vertex i : Use Path halving */
template<typename Index, typename IndexVector>
Index
etree_find(Index i, IndexVector& pp)
{
	Index p = pp(i);  // Parent
	Index gp = pp(p); // Grand parent
	while (gp != p) {
		pp(i) = gp; // Parent pointer on find path is changed to former grand parent
		i = gp;
		p = pp(i);
		gp = pp(p);
	}
	return p;
}

/** Compute the column elimination tree of a sparse matrix
 * \param mat The matrix in column-major format.
 * \param parent The elimination tree
 * \param firstRowElt The column index of the first element in each row
 * \param perm The permutation to apply to the column of \b mat
 */
template<typename MatrixType, typename IndexVector>
int
coletree(const MatrixType& mat,
		 IndexVector& parent,
		 IndexVector& firstRowElt,
		 typename MatrixType::StorageIndex* perm = 0)
{
	typedef typename MatrixType::StorageIndex StorageIndex;
	StorageIndex nc = convert_index<StorageIndex>(mat.cols()); // Number of columns
	StorageIndex m = convert_index<StorageIndex>(mat.rows());
	StorageIndex diagSize = (std::min)(nc, m);
	IndexVector root(nc); // root of subtree of etree
	root.setZero();
	IndexVector pp(nc); // disjoint sets
	pp.setZero();		// Initialize disjoint sets
	parent.resize(mat.cols());
	// Compute first nonzero column in each row
	firstRowElt.resize(m);
	firstRowElt.setConstant(nc);
	firstRowElt.segment(0, diagSize).setLinSpaced(diagSize, 0, diagSize - 1);
	bool found_diag;
	for (StorageIndex col = 0; col < nc; col++) {
		StorageIndex pcol = col;
		if (perm)
			pcol = perm[col];
		for (typename MatrixType::InnerIterator it(mat, pcol); it; ++it) {
			Index row = it.row();
			firstRowElt(row) = (std::min)(firstRowElt(row), col);
		}
	}
	/* Compute etree by Liu's algorithm for symmetric matrices,
			except use (firstRowElt[r],c) in place of an edge (r,c) of A.
	  Thus each row clique in A'*A is replaced by a star
	  centered at its first vertex, which has the same fill. */
	StorageIndex rset, cset, rroot;
	for (StorageIndex col = 0; col < nc; col++) {
		found_diag = col >= m;
		pp(col) = col;
		cset = col;
		root(cset) = col;
		parent(col) = nc;
		/* The diagonal element is treated here even if it does not exist in the matrix
		 * hence the loop is executed once more */
		StorageIndex pcol = col;
		if (perm)
			pcol = perm[col];
		for (typename MatrixType::InnerIterator it(mat, pcol); it || !found_diag;
			 ++it) { //  A sequence of interleaved find and union is performed
			Index i = col;
			if (it)
				i = it.index();
			if (i == col)
				found_diag = true;

			StorageIndex row = firstRowElt(i);
			if (row >= col)
				continue;
			rset = internal::etree_find(row, pp); // Find the name of the set containing row
			rroot = root(rset);
			if (rroot != col) {
				parent(rroot) = col;
				pp(cset) = rset;
				cset = rset;
				root(cset) = col;
			}
		}
	}
	return 0;
}

/**
 * Depth-first search from vertex n.  No recursion.
 * This routine was contributed by Cédric Doucet, CEDRAT Group, Meylan, France.
 */
template<typename IndexVector>
void
nr_etdfs(typename IndexVector::Scalar n,
		 IndexVector& parent,
		 IndexVector& first_kid,
		 IndexVector& next_kid,
		 IndexVector& post,
		 typename IndexVector::Scalar postnum)
{
	typedef typename IndexVector::Scalar StorageIndex;
	StorageIndex current = n, first, next;
	while (postnum != n) {
		// No kid for the current node
		first = first_kid(current);

		// no kid for the current node
		if (first == -1) {
			// Numbering this node because it has no kid
			post(current) = postnum++;

			// looking for the next kid
			next = next_kid(current);
			while (next == -1) {
				// No more kids : back to the parent node
				current = parent(current);
				// numbering the parent node
				post(current) = postnum++;

				// Get the next kid
				next = next_kid(current);
			}
			// stopping criterion
			if (postnum == n + 1)
				return;

			// Updating current node
			current = next;
		} else {
			current = first;
		}
	}
}

/**
 * \brief Post order a tree
 * \param n the number of nodes
 * \param parent Input tree
 * \param post postordered tree
 */
template<typename IndexVector>
void
treePostorder(typename IndexVector::Scalar n, IndexVector& parent, IndexVector& post)
{
	typedef typename IndexVector::Scalar StorageIndex;
	IndexVector first_kid, next_kid; // Linked list of children
	StorageIndex postnum;
	// Allocate storage for working arrays and results
	first_kid.resize(n + 1);
	next_kid.setZero(n + 1);
	post.setZero(n + 1);

	// Set up structure describing children
	first_kid.setConstant(-1);
	for (StorageIndex v = n - 1; v >= 0; v--) {
		StorageIndex dad = parent(v);
		next_kid(v) = first_kid(dad);
		first_kid(dad) = v;
	}

	// Depth-first search from dummy root vertex #n
	postnum = 0;
	internal::nr_etdfs(n, parent, first_kid, next_kid, post, postnum);
}

} // end namespace internal

} // end namespace Eigen

#endif // SPARSE_COLETREE_H
